#!/usr/local/bin/python3
import pyximport
# 这⾥同样指定 language_level=3
# 表示针对的是 py3
pyximport.install(language_level=3)

import read_data as rwd 
from contraction_type import *
import numpy as np
import input_output as iot
import sys 
import os
from queue import Queue
import constant as ct
import numpy  as np
import time
import threading
import argparse
import torch as th

class Task:
    def __init__(self):
        self.tasks=[]
    def registrat(self,task):
        self.tasks.append(task)
    def do(self,perams_u_th,perams_s_th,pera_dig,VdV_th,conf_id,t_se):
        for task in self.tasks:
            # t = time.localtime(time.time())
            # print(f"{type(task).__name__} begin, time: {time.asctime(t)}", flush=True)
            task.contract(perams_u_th,perams_s_th,pera_dig,VdV_th,conf_id,t_se)
            # t = time.localtime(time.time())
            # print(f"{type(task).__name__} end, time: {time.asctime(t)}", flush=True)
    def save(self):
        for task in self.tasks:
            task.save()
        
def read_data(tasks,PxPyPz,q):
    for conf_id in tasks:
        VdV_pxpypz = np.zeros((len(PxPyPz) , ct.Nt,ct.Nev1,ct.Nev1), dtype=complex)
        for t in range(ct.Nt):
            VdV_pxpypz[0,t,:,:]=np.identity(ct.Nev1) 
        for i,pxpypz in enumerate(VdV_PxPyPz[1:]):
            VdV_pxpypz[i+1,:,:,:]=iot.readin_VdV_all(ct.eigen_dir,ct.Nev1,ct.Nev1,
                    ct.Nt,conf_id,pxpypz[0],
                    pxpypz[1],pxpypz[2],ct.pad)
        

        VdV_pxpypz=VdV_pxpypz[:,:,:ct.Nev1,:ct.Nev1]
        print("reading  data begins!  ")
        t_st=time.perf_counter()
        perams_u_np = np.zeros((ct.Nt,4,ct.Nt,ct.Nev1,4,ct.Nev1),dtype=complex)
        rwd.read_peram_all_part(ct.perams_u_dirb,ct.Nt,ct.Nev,ct.Nev1,conf_id,perams_u_np)

        perams_s_np = np.zeros((ct.Nt,4,ct.Nt,ct.Nev1,4,ct.Nev1),dtype=complex)
        t_st=time.perf_counter()
        rwd.read_peram_all_part(ct.perams_s_dirb,ct.Nt,ct.Nev,ct.Nev1,conf_id,perams_s_np)
        t_ed=time.perf_counter()
        print(f"reading  data ends! t: {(t_ed-t_st)/60} ")
        

        q.put([ conf_id, perams_s_np,perams_u_np, VdV_pxpypz])
    
    q.put(-1)

dirt="/root/pion_K/data/tmp"


def contract_t_se(q):
    st  = time.perf_counter()
    while(1):
        
        perams=q.get()
        q.task_done()
        if perams==-1:
            break
        conf_id = perams[0]
        # print(f"task {conf_id}  begin!")
        perams_s_np, perams_u_np, VdV_pxpypz = perams[1:]
        perams_s_np=perams_s_np.transpose(2,0,5,3,4,1) # t_sk, t_se, e_sk,e_se,d_sk,d_se
        # perams_u_np=perams_u_np.transpose(2,0,5,3,4,1)
        perams_u_th = th.from_numpy(perams_u_np).to(device)
        perams_u_th=perams_u_th.permute(2,0,5,3,4,1)
        VdV_th =th.from_numpy(VdV_pxpypz).to(device)
        tasks=Task()
        # tasks.registrat(corr_fun_pikh(ct.X_sesk_tot0,dirt,conf_id))
        # tasks.registrat(corr_fun_pikx(ct.X_sesk_tot0,dirt,conf_id))
        # tasks.registrat(corr_fun_ki_mul(ct.P0,ct.corre_Kg_dir,["g1","g0","g2","g3"],conf_id,VdV_th))
        # tasks.registrat(corr_fun_pik_mul(ct.P_tot0,ct.corre_PiK_dir,conf_id,VdV_th))
        # tasks.registrat(corr_fun_pikkg_mul(ct.P_tot0,ct.P0,ct.corre_PiK_Kg_dir,["g1","g2","g3","g0"],conf_id,VdV_th))
        tasks.registrat(corr_fun_ki_mul(ct.P0,ct.corre_Kg_dir,["g1","g0","g2","g3"],conf_id,VdV_th))
        # tasks.registrat(corr_fun_pik_mul(ct.P_tot2,ct.corre_PiK_dir,conf_id,VdV_th))
        tasks.registrat(corr_fun_pikkg_mul(ct.P_tot0,ct.P0[0],
                                           ct.corre_PiK_Kg_dir,
                                           ["g5","g1","g2","g3"],["g1","g2","g3","g0"],conf_id,VdV_th))
        for i , pik_p in enumerate(ct.P_tot0):
            tasks.registrat(corr_fun_pik_mul_gi(pik_p,
                ct.P_tot0[i:],ct.corre_PiK_dir,conf_id,VdV_th,["g1","g5","g2"]))
        t = time.localtime(time.time())
        print(f"conf_id {conf_id}  begin, time: {time.asctime(t)}")
        _N=4
        _Nt = ct.Nt//_N
        for _n in range(0,_N):
            perams_s_th = th.from_numpy(perams_s_np[_n*_Nt:(_n+1)*_Nt]).to(device)
            for  t_sk  in range(0,_Nt):
                tasks.do(perams_u_th[t_sk],perams_s_th[t_sk],
                         VdV_th,conf_id,t_sk+_n*_Nt)
            del perams_s_th
        tasks.save()
        del tasks
        del VdV_th, perams_u_th
        del perams_s_np, perams_u_np,VdV_pxpypz
        t = time.localtime(time.time())
        print(f"conf_id {conf_id}  end, time: {time.asctime(t)}", flush=True)





if __name__=='__main__':
    parser = argparse.ArgumentParser(description='Process tasks for a HIP device') 
    parser.add_argument('--tasks', '-t', nargs='+', type=int, help='List of tasks to process')
    parser.add_argument('--init', '-i', nargs='+', type=int,
            help='initial conf_id')
    parser.add_argument('--num', '-n', nargs='+', type=int,
            help='the number of  conf')
    parser.add_argument('--step', '-s', nargs='+', type=int,
            help='the step of  conf')
    parser.add_argument('--final', '-f', nargs='+', type=int,
            help='final  conf')
    parser.add_argument('--momentum', '-m', nargs='+', type=int,
            help='total momentum')
    args = parser.parse_args()

    q=Queue(1)
    conf_ids =range(args.init[0],args.final[0],args.step[0])
    
    task_t = time.perf_counter()
    Th  =  threading.Thread(target=read_data, args=(conf_ids,VdV_PxPyPz,q) ) 
    Th.start() 
    Th1 = threading.Thread(target=contract_t_se,
            args=(q,))
    Th1.start()
    Th.join()
    Th1.join()
